import mindspore.nn as nn
import time
import os
import ctypes
import mindspore.ops as ops
from mindspore.ops import FFTWithSize
import mmap
import mindspore as ms
from mindspore import export
from mindspore import Tensor, context, dtype as mstype
import numpy as np



# 定义复数类型
class Complex(ctypes.Structure):
    _fields_ = [("real", ctypes.c_float), ("imag", ctypes.c_float)]

# 共享内存名称和大小
SHARED_MEMORY_NAME = "ZywyMem"
SHARED_MEMORY_SIZE = 250000 * ctypes.sizeof(Complex)
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
signal_ndim = 1
class ComplexFFTModel(nn.Cell):
    def __init__(self, signal_ndim):
        super(ComplexFFTModel, self).__init__()
        self.fft = FFTWithSize(signal_ndim, inverse=False, real=False)

    def construct(self, x):
        return self.fft(x)

# 定义 MainModel
class MainModel(nn.Cell):
    def __init__(self, signal_ndim):
        super(MainModel, self).__init__()
        # 实例化FFt模型
        # self.fft = FFTWithSize(signal_ndim, inverse=False, real=False)
        self.fft = ComplexFFTModel(signal_ndim)

    def construct(self, x):
        return self.fft(x)

def hello_call():
    print("v0.0.12 MindSpore version:", ms.__version__)
    # 打开共享内存
    shm_fd = open(f"/dev/shm/{SHARED_MEMORY_NAME}", "r+b")
    shm_map = mmap.mmap(shm_fd.fileno(), SHARED_MEMORY_SIZE)
    # 将共享内存映射为复数数组
    complex_array = (Complex * 250000).from_buffer(shm_map)
    real_parts = np.array([complex_struct.real for complex_struct in complex_array], dtype=np.float32)
    imag_parts = np.array([complex_struct.imag for complex_struct in complex_array], dtype=np.float32)
    numpy_complex_array = real_parts + 1j * imag_parts
    signal_ndim = 1
    model = MainModel(signal_ndim=signal_ndim)
    while True:
        if abs(complex_array[0].real - (-1.01)) > 0.001: # 如果没有新内空则睡1毫秒
            print(f'没有新内容......')
            time.sleep(0.001)
            continue
        tensor = Tensor(numpy_complex_array.astype(np.complex64), mstype.complex64)
        rst = model(tensor)
        print(f'rst: {type(rst)}; {id(rst)}; {id(numpy_complex_array)}; ??????')
        complex_array[0].real = 8.81
        complex_array[0].imag = 8.8
        time.sleep(1.5)
    shm_map.close()  # 解除内存映射
    return 'OK'

if '__main__' == __name__:
    hello_call()
